import numpy as np


def hypo(X, W):
    H = np.dot(X, W)
    return H


def error(H, Y):
    E = H - Y
    return E


def cost(E):
    m = E.shape[0]
    j = np.dot(E.T, E)[0, 0] / m
    return j


def delta_W(X, E):
    m = X.shape[0]
    dW = np.dot(X.T, E) / m
    return dW


def update_W(W, dW, alpha):
    W -= alpha * dW


def get_W_by_regular_equation(X, Y):
    W = np.dot(
        np.dot(
            np.linalg.inv(np.dot(X.T, X)),
            X.T
        ),
        Y
    )
    return W


if '__main__' == __name__:
    from python_ai.category.data.data4lin_regr import X, Y
    from mpl_toolkits import mplot3d
    import matplotlib.pyplot as plt

    np.random.seed(2)

    X_mean = X[:, 1:].mean(axis=0)
    X_sigma = X[:, 1:].std(axis=0)
    X[:, 1:] -= X_mean  # (M, 2); (2,) => (M, 2)
    X[:, 1:] /= X_sigma

    Y_mean = Y.mean(axis=0)
    Y_sigma = Y.std(axis=0)
    Y -= Y_mean
    Y /= Y_sigma

    W_re = get_W_by_regular_equation(X, Y)

    W = np.random.normal(0., 1., [X.shape[1], 1])
    # W = np.random.normal(0., 1., [X.shape[1], 1]) / X.shape[1]

    ITERS = 200
    ALPHA = 1e-2
    j_his = []
    for i in range(ITERS):
        H = hypo(X, W)
        E = error(H, Y)
        j = cost(E)
        j_his.append(j)
        dW = delta_W(X, E)
        update_W(W, dW, ALPHA)

    print('W:', W)

    plt.figure(figsize=[12, 6])
    plt_params = dict(spr=1, spc=2, spn=0)
    plt_params['spn'] += 1
    ax = plt.subplot(plt_params['spr'], plt_params['spc'], plt_params['spn'], projection='3d')
    ax.scatter3D(X[:, 1], X[:, 2], Y[:, 0], s=1)

    X_2_samples = np.array([
        [1., X[:, 1].min(), X[:, 2].min()],
        [1., X[:, 1].max(), X[:, 2].max()]
    ], dtype=np.float32)
    Y_2_samples = np.dot(X_2_samples, W)
    Y_2_samples_re = np.dot(X_2_samples, W_re)
    ax.plot3D(X_2_samples[:, 1], X_2_samples[:, 2], Y_2_samples[:, 0], color='red')
    ax.plot3D(X_2_samples[:, 1], X_2_samples[:, 2], Y_2_samples_re[:, 0], color='black')

    plt_params['spn'] += 1
    plt.subplot(plt_params['spr'], plt_params['spc'], plt_params['spn'])
    plt.title('Cost history')
    plt.plot(j_his)
    plt.show()
